package com.bff.gaia.runtime.maqy;

import com.bff.gaia.api.common.functions.RichMapPartitionFunction;
import com.bff.gaia.api.java.tuple.Tuple2;
import com.bff.gaia.configuration.ConfigConstants;
import com.bff.gaia.configuration.Configuration;
import com.bff.gaia.configuration.GlobalConfiguration;
import com.bff.gaia.runtime.maqy.service.Rate;
import com.bff.gaia.util.Collector;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Iterator;

/**
 * 此Function用于统计这个子任务中的数据量（记录数）
 */
public class AttachPartition<T> extends RichMapPartitionFunction<T, Tuple2<HostInfo, Long>> {
	private int subTaskIdx = -1;

	@Override
	public void mapPartition(Iterable<T> values, Collector<Tuple2<HostInfo, Long>> out) throws Exception {
		if (subTaskIdx == -1) {
			throw new Exception("SubTaskIdx must be set!");
		}

		String hostName = "localhost";
		String ip = "127.0.0.1";
		try {
			hostName = InetAddress.getLocalHost().getHostName();
			ip = InetAddress.getLocalHost().getHostAddress();
		} catch (UnknownHostException e) {
			e.printStackTrace();
		}

		// test
		// Rate testRate = NetUtil.getCurrentRate();
		// Log.info("hfy", String.format("[AttachPartition] Current rate: %s", testRate));

		Configuration configuration = GlobalConfiguration.loadConfiguration();
		double upLink = configuration.getDouble(ConfigConstants.UPLINK_BANDWIDTH, -1);
		double downLink = configuration.getDouble(ConfigConstants.DOWNLINK_BANDWIDTH, -1);
		Rate rate = new Rate(upLink, downLink);

		HostInfo key = new HostInfo(hostName, ip, rate, this.subTaskIdx);

		long count = 0;
		Iterator iterator = values.iterator();
		while (iterator.hasNext()) {
			count++;
			iterator.next();
		}
		out.collect(new Tuple2<>(key, count));
	}

	public void setSubTaskIdx(int subTaskIdx) {
		this.subTaskIdx = subTaskIdx;
	}
}